package com.example.gpu;

import jcuda.*;
import jcuda.driver.*;

import static jcuda.driver.JCudaDriver.*;

public class ImprovedJCudaVectorAddExample {
    
    // CUDA核函数代码
    private static final String CUDA_CODE = 
        "extern \"C\" \n" +
        "__global__ void add(int n, float *a, float *b, float *c)\n" +
        "{\n" +
        "    int i = blockIdx.x * blockDim.x + threadIdx.x;\n" +
        "    if (i < n) {\n" +
        "        c[i] = a[i] + b[i];\n" +
        "    }\n" +
        "}\n";

    public static void main(String[] args) {
        // 使用工具类初始化CUDA
        CUcontext context = CudaUtils.initCUDA(0);
        if (context == null) {
            System.out.println("CUDA初始化失败，无法继续");
            return;
        }
        
        // 向量大小
        int n = 1000000;
        
        // 为输入和输出数据分配内存
        float[] a = new float[n];
        float[] b = new float[n];
        float[] c = new float[n];
        
        // 初始化输入数据
        for (int i = 0; i < n; i++) {
            a[i] = (float)i;
            b[i] = (float)i * 2.0f;
        }
        
        try {
            // 在GPU上分配内存并复制数据
            CUdeviceptr d_a = CudaUtils.copyToDevice(a);
            CUdeviceptr d_b = CudaUtils.copyToDevice(b);
            CUdeviceptr d_c = new CUdeviceptr();
            cuMemAlloc(d_c, n * Sizeof.FLOAT);
            
            // 设置线程块和网格大小
            int blockSize = 256;
            int gridSize = CudaUtils.calculateGridSize(n, blockSize);
            
            // 这里我们使用简化的方式来加载核函数
            // 在实际应用中，应该使用CudaUtils.loadKernelFromPTX方法加载编译好的PTX文件
            CUmodule module = new CUmodule();
            CUfunction function = new CUfunction();
            
            System.out.println("注意: 要运行此示例，需要将CUDA代码编译为PTX文件");
            System.out.println("然后修改代码使用CudaUtils.loadKernelFromPTX加载编译好的文件");
            
            // 模拟核函数加载，实际运行时需要替换为真实的加载代码
            // Object[] kernelObj = CudaUtils.loadKernelFromPTX("vectorAdd.ptx", "add");
            // if (kernelObj == null) {
            //     System.out.println("加载核函数失败");
            //     return;
            // }
            // module = (CUmodule)kernelObj[0];
            // function = (CUfunction)kernelObj[1];
            
            // 由于我们无法直接加载核函数，这里只进行CPU计算
            System.out.println("进行CPU计算...");
            long startTime = System.currentTimeMillis();
            for (int i = 0; i < n; i++) {
                c[i] = a[i] + b[i];
            }
            long endTime = System.currentTimeMillis();
            
            System.out.println("CPU计算耗时: " + (endTime - startTime) + " ms");
            System.out.println("前10个计算结果: ");
            for (int i = 0; i < Math.min(10, n); i++) {
                System.out.printf("c[%d] = %.1f\n", i, c[i]);
            }
            
            // 释放GPU资源
            CudaUtils.releaseCUDA(null, module, d_a, d_b, d_c);
            
        } catch (Exception e) {
            System.err.println("执行GPU计算时出错: " + e.getMessage());
            e.printStackTrace();
        } finally {
            // 确保释放上下文
            if (context != null) {
                cuCtxDestroy(context);
            }
        }
    }
    
    /**
     * 将CUDA代码保存到文件
     * @param code CUDA代码
     * @param filePath 输出文件路径
     */
    private static void saveCudaCodeToFile(String code, String filePath) {
        try {
            java.io.FileWriter writer = new java.io.FileWriter(filePath);
            writer.write(code);
            writer.close();
            System.out.println("CUDA代码已保存到: " + filePath);
        } catch (java.io.IOException e) {
            System.err.println("保存CUDA代码失败: " + e.getMessage());
        }
    }
}